package com.jhm.ai.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;

/**
 * 简单的sse请求工具
 */
public class SseRequest {
    private static final Logger log = LoggerFactory.getLogger(SseRequest.class);

    @FunctionalInterface
    public interface SseHandler {
        void handle(String line, PrintWriter writer) throws IOException;
    }

    public static void connect(String url, String param, HttpServletResponse response, SseHandler handler) throws IOException {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");

        PrintWriter writer = response.getWriter();
        BufferedReader reader = null;
        HttpURLConnection conn = null;

        try {
            URL realUrl = new URL(url);
            conn = (HttpURLConnection) realUrl.openConnection();
            conn.setRequestMethod("POST");
            conn.setRequestProperty("Content-Type", "application/json");
            conn.setRequestProperty("Accept", "text/event-stream");
            conn.setDoOutput(true);
            conn.setDoInput(true);
            conn.setConnectTimeout(30000);
            conn.setReadTimeout(0);

            try (PrintWriter out = new PrintWriter(conn.getOutputStream())) {
                out.print(param);
                out.flush();
            }

            // 检查响应码
            int responseCode = conn.getResponseCode();
            if (responseCode != HttpURLConnection.HTTP_OK) {
                // 读取错误流
                try (BufferedReader errorReader = new BufferedReader(
                        new InputStreamReader(conn.getErrorStream(), StandardCharsets.UTF_8))) {
                    StringBuilder errorResponse = new StringBuilder();
                    String line;
                    while ((line = errorReader.readLine()) != null) {
                        errorResponse.append(line);
                    }
                    log.error("Server returned error code: {} with message: {}", responseCode, errorResponse.toString());
                    
                    // 发送错误消息给客户端
                    writer.write("data: {\"error\": \"Server error: " + responseCode + "\", \"type\": \"error\"}\n\n");
                    writer.flush();
                    return;
                }
            }

            reader = new BufferedReader(new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8));
            String line;
            
            while ((line = reader.readLine()) != null) {
                if (line.trim().isEmpty()) continue;
                if (handler != null) {
                    handler.handle(line, writer);
                }
            }
        } catch (Exception e) {
            log.error("SSE request error", e);
            String errorMessage;
            if (e instanceof java.net.SocketTimeoutException) {
                errorMessage = "服务器响应超时，请稍后重试";
            } else if (e instanceof java.net.ConnectException) {
                errorMessage = "无法连接到服务器，请检查网络";
            } else {
                errorMessage = "服务器繁忙，请稍后重试";
            }
            
            String jsonError = String.format("data: {\"error\": \"%s\", \"type\": \"error\"}\n\n", 
                errorMessage.replace("\"", "'"));
            writer.write(jsonError);
            writer.flush();
        } finally {
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e) {
                    log.error("Error closing reader", e);
                }
            }
            if (conn != null) {
                conn.disconnect();
            }
            writer.close();
        }
    }
}
